import pygame
import random
import sys

# 初始化pygame
pygame.init()

# 游戏常量
WINDOW_WIDTH = 800
WINDOW_HEIGHT = 600
GRID_SIZE = 20
CELL_SIZE = 20
FPS = 10

# 颜色定义
BLACK = (0, 0, 0)
WHITE = (255, 255, 255)
GREEN = (0, 255, 0)
RED = (255, 0, 0)
GRAY = (128, 128, 128)

class Snake:
    def __init__(self):
        self.positions = [(200, 200), (180, 200), (160, 200)]  # 初始3节
        self.direction = (20, 0)  # 向右移动
        self.grow = False
        
    def move(self):
        head_x, head_y = self.positions[0]
        dx, dy = self.direction
        
        new_head = (head_x + dx, head_y + dy)
        self.positions.insert(0, new_head)
        
        if not self.grow:
            self.positions.pop()
        else:
            self.grow = False
            
    def change_direction(self, new_direction):
        # 防止反向移动
        current_dx, current_dy = self.direction
        new_dx, new_dy = new_direction
        
        if (current_dx, current_dy) != (-new_dx, -new_dy):
            self.direction = new_direction
            
    def check_collision(self):
        head = self.positions[0]
        
        # 检查墙壁碰撞
        if (head[0] < 0 or head[0] >= WINDOW_WIDTH or
            head[1] < 0 or head[1] >= WINDOW_HEIGHT):
            return True
            
        # 检查自身碰撞
        if head in self.positions[1:]:
            return True
            
        return False
        
    def grow_snake(self):
        self.grow = True

class Food:
    def __init__(self):
        self.position = self.generate_position()
        
    def generate_position(self):
        x = random.randint(0, (WINDOW_WIDTH - CELL_SIZE) // CELL_SIZE) * CELL_SIZE
        y = random.randint(0, (WINDOW_HEIGHT - CELL_SIZE) // CELL_SIZE) * CELL_SIZE
        return (x, y)
        
    def respawn(self, snake_positions):
        while True:
            self.position = self.generate_position()
            if self.position not in snake_positions:
                break

class Game:
    def __init__(self):
        self.screen = pygame.display.set_mode((WINDOW_WIDTH, WINDOW_HEIGHT))
        pygame.display.set_caption("贪吃蛇游戏")
        self.clock = pygame.time.Clock()
        self.font = pygame.font.Font(None, 36)
        
        self.reset_game()
        
    def reset_game(self):
        self.snake = Snake()
        self.food = Food()
        self.score = 0
        self.game_over = False
        self.paused = False
        
    def handle_events(self):
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return False
                
            if event.type == pygame.KEYDOWN:
                if event.key == pygame.K_ESCAPE:
                    return False
                    
                if not self.game_over and not self.paused:
                    if event.key == pygame.K_UP:
                        self.snake.change_direction((0, -20))
                    elif event.key == pygame.K_DOWN:
                        self.snake.change_direction((0, 20))
                    elif event.key == pygame.K_LEFT:
                        self.snake.change_direction((-20, 0))
                    elif event.key == pygame.K_RIGHT:
                        self.snake.change_direction((20, 0))
                        
                if event.key == pygame.K_SPACE:
                    if self.game_over:
                        self.reset_game()
                    else:
                        self.paused = not self.paused
                        
        return True
        
    def update(self):
        if not self.game_over and not self.paused:
            self.snake.move()
            
            # 检查是否吃到食物
            if self.snake.positions[0] == self.food.position:
                self.score += 10
                self.snake.grow_snake()
                self.food.respawn(self.snake.positions)
                
            # 检查碰撞
            if self.snake.check_collision():
                self.game_over = True
                
    def draw(self):
        self.screen.fill(BLACK)
        
        # 绘制网格
        for x in range(0, WINDOW_WIDTH, CELL_SIZE):
            pygame.draw.line(self.screen, GRAY, (x, 0), (x, WINDOW_HEIGHT))
        for y in range(0, WINDOW_HEIGHT, CELL_SIZE):
            pygame.draw.line(self.screen, GRAY, (0, y), (WINDOW_WIDTH, y))
            
        # 绘制蛇
        for position in self.snake.positions:
            pygame.draw.rect(self.screen, GREEN, 
                           (position[0], position[1], CELL_SIZE, CELL_SIZE))
            
        # 绘制食物
        pygame.draw.rect(self.screen, RED,
                       (self.food.position[0], self.food.position[1], 
                        CELL_SIZE, CELL_SIZE))
        
        # 绘制分数
        score_text = self.font.render(f"分数: {self.score}", True, WHITE)
        self.screen.blit(score_text, (10, 10))
        
        # 绘制游戏状态
        if self.paused:
            pause_text = self.font.render("游戏暂停 - 按空格键继续", True, WHITE)
            text_rect = pause_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2))
            self.screen.blit(pause_text, text_rect)
            
        if self.game_over:
            game_over_text = self.font.render("游戏结束!", True, WHITE)
            score_text = self.font.render(f"最终分数: {self.score}", True, WHITE)
            restart_text = self.font.render("按空格键重新开始", True, WHITE)
            
            game_over_rect = game_over_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2 - 50))
            score_rect = score_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2))
            restart_rect = restart_text.get_rect(center=(WINDOW_WIDTH//2, WINDOW_HEIGHT//2 + 50))
            
            self.screen.blit(game_over_text, game_over_rect)
            self.screen.blit(score_text, score_rect)
            self.screen.blit(restart_text, restart_rect)
            
        pygame.display.flip()
        
    def run(self):
        running = True
        while running:
            running = self.handle_events()
            self.update()
            self.draw()
            self.clock.tick(FPS)
            
        pygame.quit()
        sys.exit()

if __name__ == "__main__":
    game = Game()
    game.run()
